import pickle
import os

import numpy as np
import torch

class CIFAR(object):
    '''
    folder     : 数据集存放路径
    name       : 数据文件（二进制）名称
    list_file  : 数据集序列文件名称
    '''
    def __init__(self, kwargs, transforms=None):
        super().__init__()
        self.folder = kwargs['folder']
        self.name = kwargs['name']
        self.path = os.path.join(self.folder, self.name)
        self.transfoms = transforms
        if 'list_file' in kwargs:
            self.list_path = os.path.join(self.folder, kwargs['list_file'])
        else:
            self.list_path = None

        fp = open(self.path, 'rb')
        self.contents = pickle.load(fp, encoding='latin1')
        fp.close()

        if self.list_path != None:
            fp = open(self.list_path)
            self.ids = []
            for line in fp.readlines():
                line = line.strip('\n').strip()
                line_ = line.split(' ')
                self.ids.append(int(line_[0]))
            fp.close()

    def __getitem__(self, index):
        if self.list_path == None:
            ids = index
        else:
            ids = int(self.ids[index])

        dict = {
            'filenames' : self.contents['filenames'][ids],
            'fine_labels' : self.contents['fine_labels'][ids],
            'coarse_labels' : self.contents['coarse_labels'][ids],
            'data' : self.contents['data'][ids].reshape(3, 32, 32),
        }

        if self.transfoms != None:
            dict['data'] = self.transfoms(dict['data'].transpose(1, 2, 0))

        return dict

    def __len__(self):
        if self.list_path == None:
            return len(self.contents['filenames'])
        else:
            return len(self.ids)

if __name__ == '__main__':
    import matplotlib.pyplot as plt
    import torchvision

    src_path = '../CIFAR/CIFAR100'
    tgt_path = '../CIFAR/CIFAR100'

    kwargs = {
        'folder'     : src_path,
        'name'       : 'train',
        # 'name'       : 'train_list',
        # 'list_file'  : 'train_list.txt',
    }

    dataset_trian = CIFAR(
        kwargs, 
        torchvision.transforms.Compose(
            [
                torchvision.transforms.ToTensor(),
                torchvision.transforms.Resize(224),
            ]
        )
    )
    for data in dataset_trian:
        np_photo = data['data'].numpy()
        np_photo = np_photo.transpose(1, 2, 0)
        
        plt.imshow(np_photo)
        plt.show()
        

